#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
IR 字面量表达式节点定义
"""

from dataclasses import dataclass, field
from typing import List as TypingList, Iterator
import ast
from .base import IRExpr, build_expr_from_ast
from ..base import IRNode, get_default_ir_node


@dataclass
class IRDict(IRExpr):
    """字典字面量"""

    keys: TypingList[IRNode] = field(default_factory=list)
    values: TypingList[IRNode] = field(default_factory=list)
    node_type: str = "dict"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        for key in self.keys:
            yield key
        for value in self.values:
            yield value
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.Dict) -> None:
        """从 AST 字典节点构建 IR 字典节点内容

        Args:
            node: AST 字典节点
        """
        super().build_from_ast(node)

        self.keys = []
        self.values = []

        for key, value in zip(node.keys, node.values):
            if key is not None:  # 字典解包时 key 可能为 None
                key_ir = build_expr_from_ast(key, self)
                self.keys.append(key_ir)
            else:
                # 处理字典解包 {**other_dict}
                self.keys.append(get_default_ir_node())

            value_ir = build_expr_from_ast(value, self)
            self.values.append(value_ir)


@dataclass
class IRSet(IRExpr):
    """集合字面量"""

    elements: TypingList[IRNode] = field(default_factory=list)
    node_type: str = "set"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        for element in self.elements:
            yield element
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.Set) -> None:
        """从 AST 集合节点构建 IR 集合节点内容

        Args:
            node: AST 集合节点
        """
        super().build_from_ast(node)
        self.elements = []
        for element in node.elts:
            element_ir = build_expr_from_ast(element, self)
            self.elements.append(element_ir)


@dataclass
class IRTuple(IRExpr):
    """元组字面量"""

    elements: TypingList[IRNode] = field(default_factory=list)
    node_type: str = "tuple"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        for element in self.elements:
            yield element
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.Tuple) -> None:
        """从 AST 元组节点构建 IR 元组节点内容

        Args:
            node: AST 元组节点
        """
        super().build_from_ast(node)
        self.elements = []
        for element in node.elts:
            element_ir = build_expr_from_ast(element, self)
            self.elements.append(element_ir)


@dataclass
class IRList(IRExpr):
    """列表字面量"""

    elements: TypingList[IRNode] = field(default_factory=list)
    node_type: str = "list"

    def get_child_nodes(self) -> Iterator[IRNode]:
        """获取所有子节点"""
        for element in self.elements:
            yield element
        yield from super().get_child_nodes()

    def build_from_ast(self, node: ast.List) -> None:
        """从 AST 列表节点构建 IR 列表节点内容

        Args:
            node: AST 列表节点
        """
        super().build_from_ast(node)
        self.elements = []
        for element in node.elts:
            element_ir = build_expr_from_ast(element, self)
            self.elements.append(element_ir)
